{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Memory in LlamaIndex\n",
    "\n",
    "The `Memory` class in LlamaIndex is used to store and retrieve both short-term and long-term memory.\n",
    "\n",
    "You can use it on its own and orchestrate within a custom workflow, or use it within an existing agent.\n",
    "\n",
    "By default, short-term memory is represented as a FIFO queue of `ChatMessage` objects. Once the queue exceeds a certain size, the last X messages within a flush size are archived and optionally flushed to long-term memory blocks.\n",
    "\n",
    "Long-term memory is represented as `Memory Block` objects. These objects receive the messages that are flushed from short-term memory, and optionally process them to extract information. Then when memory is retrieved, the short-term and long-term memories are merged together."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "This notebook will use `OpenAI` as an LLM/embedding model for various parts of the example.\n",
    "\n",
    "For vector retrieval, we will rely on `Chroma` as a vector store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-core llama-index-llms-openai llama-index-embeddings-openai llama-index-vector-stores-chroma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-proj-...\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Short-term Memory\n",
    "\n",
    "Let's explore how to configure various components of short-term memory.\n",
    "\n",
    "For visual purposes, we will set some low token limits to more easily observe the memory behavior.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.memory import Memory\n",
    "\n",
    "memory = Memory.from_defaults(\n",
    "    session_id=\"my_session\",\n",
    "    token_limit=50,  # Normally you would set this to be closer to the LLM context window (i.e. 75,000, etc.)\n",
    "    token_flush_size=10,\n",
    "    chat_history_token_ratio=0.7,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's review the configuration we used and what it means:\n",
    "\n",
    "- `session_id`: A unique identifier for the session. Used to mark chat messages in a SQL database as belonging to a specific session.\n",
    "- `token_limit`: The maximum number of tokens that can be stored in short-term + long-term memory.\n",
    "- `chat_history_token_ratio`: The ratio of tokens in the short-term chat history to the total token limit. Here this means that 50*0.7 = 35 tokens are allocated to short-term memory, and the rest is allocated to long-term memory.\n",
    "- `token_flush_size`: The number of tokens to flush to long-term memory when the token limit is exceeded. Note that we did not configure long-term memory, so these messages are merely archived in the database and removed from the short-term memory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using our memory, we can manually add some messages and observe how it works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "# Simulate a long conversation\n",
    "for i in range(100):\n",
    "    await memory.aput_messages(\n",
    "        [\n",
    "            ChatMessage(role=\"user\", content=\"Hello, world!\"),\n",
    "            ChatMessage(role=\"assistant\", content=\"Hello, world to you too!\"),\n",
    "            ChatMessage(role=\"user\", content=\"What is the capital of France?\"),\n",
    "            ChatMessage(\n",
    "                role=\"assistant\", content=\"The capital of France is Paris.\"\n",
    "            ),\n",
    "        ]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since our token limit is small, we will only see the last 4 messages in short-term memory (since this fits withint the 50*0.7 limit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "user: Hello, world!\n",
      "assistant: Hello, world to you too!\n",
      "user: What is the capital of France?\n",
      "assistant: The capital of France is Paris.\n"
     ]
    }
   ],
   "source": [
    "current_chat_history = await memory.aget()\n",
    "for msg in current_chat_history:\n",
    "    print(msg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we retrieva all messages, we will find all 400 messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "400\n"
     ]
    }
   ],
   "source": [
    "all_messages = await memory.aget_all()\n",
    "print(len(all_messages))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can clear the memory at any time to start fresh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "await memory.areset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n"
     ]
    }
   ],
   "source": [
    "all_messages = await memory.aget_all()\n",
    "print(len(all_messages))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Long-term Memory\n",
    "\n",
    "Long-term memory is represented as `Memory Block` objects. These objects receive the messages that are flushed from short-term memory, and optionally process them to extract information. Then when memory is retrieved, the short-term and long-term memories are merged together.\n",
    "\n",
    "LlamaIndex provides 3 prebuilt memory blocks:\n",
    "\n",
    "- `StaticMemoryBlock`: A memory block that stores a static piece of information.\n",
    "- `FactExtractionMemoryBlock`: A memory block that extracts facts from the chat history.\n",
    "- `VectorMemoryBlock`: A memory block that stores and retrieves batches of chat messages from a vector database.\n",
    "\n",
    "Each block has a `priority` that is used when the long-term memory + short-term memory exceeds the token limit. Priority 0 means the block will always be kept in memory, priority 1 means the block will be temporarily disabled, and so on. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.memory import (\n",
    "    StaticMemoryBlock,\n",
    "    FactExtractionMemoryBlock,\n",
    "    VectorMemoryBlock,\n",
    ")\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.vector_stores.chroma import ChromaVectorStore\n",
    "import chromadb\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4.1-mini\")\n",
    "embed_model = OpenAIEmbedding(model=\"text-embedding-3-small\")\n",
    "\n",
    "client = chromadb.EphemeralClient()\n",
    "vector_store = ChromaVectorStore(\n",
    "    chroma_collection=client.get_or_create_collection(\"test_collection\")\n",
    ")\n",
    "\n",
    "blocks = [\n",
    "    StaticMemoryBlock(\n",
    "        name=\"core_info\",\n",
    "        static_content=\"My name is Logan, and I live in Saskatoon. I work at LlamaIndex.\",\n",
    "        priority=0,\n",
    "    ),\n",
    "    FactExtractionMemoryBlock(\n",
    "        name=\"extracted_info\",\n",
    "        llm=llm,\n",
    "        max_facts=50,\n",
    "        priority=1,\n",
    "    ),\n",
    "    VectorMemoryBlock(\n",
    "        name=\"vector_memory\",\n",
    "        # required: pass in a vector store like qdrant, chroma, weaviate, milvus, etc.\n",
    "        vector_store=vector_store,\n",
    "        priority=2,\n",
    "        embed_model=embed_model,\n",
    "        # The top-k message batches to retrieve\n",
    "        # similarity_top_k=2,\n",
    "        # optional: How many previous messages to include in the retrieval query\n",
    "        # retrieval_context_window=5\n",
    "        # optional: pass optional node-postprocessors for things like similarity threshold, etc.\n",
    "        # node_postprocessors=[...],\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With our blocks created, we can pass them into the `Memory` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.memory import Memory\n",
    "\n",
    "memory = Memory.from_defaults(\n",
    "    session_id=\"my_session\",\n",
    "    token_limit=30000,\n",
    "    # Setting a extremely low ratio so that more tokens are flushed to long-term memory\n",
    "    chat_history_token_ratio=0.02,\n",
    "    token_flush_size=500,\n",
    "    memory_blocks=blocks,\n",
    "    # insert into the latest user message, can also be \"system\"\n",
    "    insert_method=\"user\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this, we can simulate a conversation with an agent and inspect the long-term memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    tools=[],\n",
    "    llm=llm,\n",
    ")\n",
    "\n",
    "user_msgs = [\n",
    "    \"Hi! My name is Logan\",\n",
    "    \"What is your opinion on minature shnauzers?\",\n",
    "    \"Do they shed a lot?\",\n",
    "    \"What breeds are comparable in size?\",\n",
    "    \"What is your favorite breed?\",\n",
    "    \"Would you recommend owning a dog?\",\n",
    "    \"What should I buy to prepare for owning a dog?\",\n",
    "]\n",
    "\n",
    "for user_msg in user_msgs:\n",
    "    _ = await agent.run(user_msg=user_msg, memory=memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's inspect the most recent user-message and see what the memory inserts into the user message.\n",
    "\n",
    "Note that we pass in at least one chat message so that the vector memory actually runs retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_history = await memory.aget()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n"
     ]
    }
   ],
   "source": [
    "print(len(chat_history))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great, we can see that the current FIFO queue is only 2 messages (expected since we set the chat history token ratio to 0.02).\n",
    "\n",
    "Now, let's inspect the long-term memory blocks that are inserted into the latest user message."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<memory>\n",
      "<core_info>\n",
      "My name is Logan, and I live in Saskatoon. I work at LlamaIndex.\n",
      "</core_info>\n",
      "<extracted_info>\n",
      "<fact>User's name is Logan</fact>\n",
      "<fact>User lives in Saskatoon</fact>\n",
      "<fact>User works at LlamaIndex</fact>\n",
      "<fact>User is interested in Miniature Schnauzers</fact>\n",
      "</extracted_info>\n",
      "<vector_memory>\n",
      "<message role='user'>Hi! My name is Logan</message>\n",
      "<message role='assistant'>Hi Logan! Nice to meet you. How can I assist you today?</message>\n",
      "<message role='user'>What is your opinion on minature shnauzers?</message>\n",
      "<message role='assistant'>Hi Logan! Miniature Schnauzers are wonderful dogs—they're known for being intelligent, friendly, and energetic. They often make great companions because they're loyal and good with families. Plus, their distinctive beard and eyebrows give them a charming, expressive look. Do you have one, or are you thinking about getting one?</message>\n",
      "<message role='user'>Do they shed a lot?</message>\n",
      "<message role='assistant'>Hi Logan! Miniature Schnauzers are actually known for being low shedders. They have a wiry double coat that doesn't shed much, which makes them a good choice for people who prefer a cleaner home or have mild allergies. However, their coat does require regular grooming and trimming to keep it looking its best. Since you’re in Saskatoon, the grooming routine might also help keep them comfortable through the changing seasons. Are you considering getting one as a pet?</message>\n",
      "<message role='user'>What breeds are comparable in size?</message>\n",
      "<message role='assistant'>Hi Logan! Miniature Schnauzers typically weigh between 11 to 20 pounds (5 to 9 kg) and stand about 12 to 14 inches (30 to 36 cm) tall at the shoulder. Breeds comparable in size include:\n",
      "\n",
      "- **Cairn Terrier**  \n",
      "- **West Highland White Terrier (Westie)**  \n",
      "- **Scottish Terrier**  \n",
      "- **Pomeranian** (though usually a bit smaller)  \n",
      "- **Beagle** (on the smaller side of the breed)  \n",
      "- **French Bulldog** (a bit stockier but similar in height)  \n",
      "\n",
      "These breeds are similar in size and can have comparable energy levels and grooming needs, depending on the breed. If you’re thinking about a dog that fits well with your lifestyle in Saskatoon and your work at LlamaIndex, I’d be happy to help you explore options!</message>\n",
      "<message role='user'>What is your favorite breed?</message>\n",
      "<message role='assistant'>Hi Logan! I don't have personal preferences, but I really appreciate breeds like the Miniature Schnauzer because of their intelligence, friendly nature, and low-shedding coat. They seem like great companions, especially for someone living in a place with changing seasons like Saskatoon. Do you have a favorite breed, or one you’re particularly interested in?</message>\n",
      "<message role='user'>Would you recommend owning a dog?</message>\n",
      "<message role='assistant'>Hi Logan! Owning a dog can be a wonderful experience, offering companionship, exercise, and even stress relief. Since you live in Saskatoon, where the seasons can be quite distinct, a dog can be a great motivator to get outside and enjoy the fresh air year-round.\n",
      "\n",
      "That said, it’s important to consider your lifestyle and work schedule at LlamaIndex. Dogs require time, attention, and care—regular walks, playtime, grooming, and vet visits. If you have the time and energy to commit, a dog can be a fantastic addition to your life. Breeds like Miniature Schnauzers, which are adaptable and relatively low-maintenance in terms of shedding, might be a good fit.\n",
      "\n",
      "If you’re unsure, maybe start by volunteering at a local animal shelter or fostering a dog to see how it fits with your routine. Would you like tips on how to prepare for dog ownership or suggestions on breeds that suit your lifestyle?</message>\n",
      "</vector_memory>\n",
      "</memory>\n",
      "What should I buy to prepare for owning a dog?\n"
     ]
    }
   ],
   "source": [
    "for block in chat_history[-2].blocks:\n",
    "    print(block.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use this memory outside an agent, and to highlight more of the usage, you might do something like the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You were asking about Miniature Schnauzers.\n"
     ]
    }
   ],
   "source": [
    "new_user_msg = ChatMessage(\n",
    "    role=\"user\", content=\"What kind of dog was I asking about?\"\n",
    ")\n",
    "await memory.aput(new_user_msg)\n",
    "\n",
    "# Get the new chat history\n",
    "new_chat_history = await memory.aget()\n",
    "resp = await llm.achat(new_chat_history)\n",
    "await memory.aput(resp.message)\n",
    "print(resp.message.content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
